# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         ba_solver.py
# Date:           2025-05-24
# Description:    本模块实现了基于 Ceres Solver 的非线性优化算法，用于对多视角下的位姿进行联合优化（Bundle Adjustment）。
#                 主要功能包括：
#                 - 定义重投影误差函数（ReprojectionError）；
#                 - 构建 Bundle Adjustment 优化问题；
#                 - 使用 Levenberg-Marquardt 算法优化旋转和平移向量；
#                 是整个骰子姿态识别系统中视觉几何优化与多帧一致性提升的核心模块。
# ==============================================================================

import numpy as np
import pyceres as ceres
import cv2


class ReprojectionError(ceres.CostFunction):
    def __init__(self, points_2d, fx, fy, cx, cy, points_3d):
        super().__init__()
        self.points_2d = np.array(points_2d, dtype=np.float64)
        self.fx = fx
        self.fy = fy
        self.cx = cx
        self.cy = cy
        self.points_3d = np.array(points_3d, dtype=np.float64)

        self.set_num_residuals(2 * len(points_3d))
        self.set_parameter_block_sizes([3, 3])

    def Evaluate(self, parameters, residuals, jacobians):
        rvec = parameters[0].copy()
        tvec = parameters[1].copy()

        R = cv2.Rodrigues(rvec)[0]

        for i, point_3d in enumerate(self.points_3d):
            # 投影到相机坐标系
            point_camera = R @ point_3d + tvec

            if point_camera[2] == 0:
                residuals[2 * i] = 0.0
                residuals[2 * i + 1] = 0.0
                continue

            # 归一化坐标
            x = point_camera[0] / point_camera[2]
            y = point_camera[1] / point_camera[2]

            predicted_x = self.fx * x + self.cx
            predicted_y = self.fy * y + self.cy

            # 残差
            residuals[2 * i] = self.points_2d[i][0] - predicted_x
            residuals[2 * i + 1] = self.points_2d[i][1] - predicted_y

        if jacobians is None:
            return True

        if jacobians[0] is not None or jacobians[1] is not None:
            self._numerical_jacobian(parameters, residuals, jacobians)

        return True

    def _numerical_jacobian(self, parameters, residuals, jacobians, eps=1e-5):
        num_points = len(self.points_3d)
        num_residuals = 2 * num_points

        original_params = [p.copy() for p in parameters]

        original_residuals = residuals.copy()

        for block_idx in range(len(original_params)):
            param_block = original_params[block_idx]
            jac = jacobians[block_idx]

            if jac is None:
                continue

            param_dim = len(param_block)

            for i in range(param_dim):
                # 正向扰动
                param_block[i] += eps
                pos_residuals = np.zeros(num_residuals)
                self.Evaluate(original_params, pos_residuals, None)

                # 负向扰动
                param_block[i] -= 2 * eps
                neg_residuals = np.zeros(num_residuals)
                self.Evaluate(original_params, neg_residuals, None)

                param_block[i] += eps

                derivative = (pos_residuals - neg_residuals) / (2 * eps)

                for res_idx in range(num_residuals):
                    idx = res_idx * param_dim + i
                    jac[idx] = derivative[res_idx]

        # 恢复原始参数和残差
        for p, orig in zip(parameters, original_params):
            p[:] = orig[:]
        residuals[:] = original_residuals[:]

        return True


def bundle_adjustment(corners_list, rvec_list, tvec_list, huber_loss):
    from configs.norm_data import fx, fy, cx, cy, dice_3d_points
    from configs.config import Config

    points_3d = dice_3d_points[4:8]

    problem = ceres.Problem()

    opt_params_list = []
    for i in range(len(corners_list)):
        rvec = rvec_list[i].copy().ravel().astype(np.float64)
        tvec = tvec_list[i].copy().ravel().astype(np.float64)
        params = [rvec, tvec]
        opt_params_list.append(params)

        cost_function = ReprojectionError(
            corners_list[i], fx, fy, cx, cy, points_3d
        )
        loss_function = ceres.HuberLoss(Config.HUBER_LOSS_VALUE)
        problem.add_residual_block(cost_function, loss_function, params)

    options = ceres.SolverOptions()
    options.linear_solver_type = ceres.LinearSolverType.DENSE_QR
    options.minimizer_type = ceres.MinimizerType.TRUST_REGION
    options.trust_region_strategy_type = ceres.TrustRegionStrategyType.LEVENBERG_MARQUARDT
    options.max_num_iterations = 100
    options.function_tolerance = 1e-10
    options.gradient_tolerance = 1e-10
    options.parameter_tolerance = 1e-10

    summary = ceres.SolverSummary()
    ceres.solve(options, problem, summary)

    opt_rvecs = [params[0].copy() for params in opt_params_list]
    opt_tvecs = [params[1].copy() for params in opt_params_list]

    return opt_rvecs, opt_tvecs